package com.vova.rec.tensorflow.config;

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import tensorflow.serving.ModelServiceGrpc;
import tensorflow.serving.PredictionServiceGrpc;
import tensorflow.serving.SessionServiceGrpc;

/**
 * Java client for access Tensorflow Serving
 */
@Configuration
public class GrpcClientConfig {
    @Bean
    public ManagedChannel managedChannel(ModelServerProperties modelServerProperties) {
        return ManagedChannelBuilder.forAddress(modelServerProperties.getHost(), modelServerProperties.getPort())
                .usePlaintext()
                .build();
    }

    @Bean
    public PredictionServiceGrpc.PredictionServiceBlockingStub predictionServiceStub(ManagedChannel managedChannel) {
        return PredictionServiceGrpc.newBlockingStub(managedChannel);
    }

    /**
     * 获取当前模型状态，更新模型配置
     */
    @Bean
    public ModelServiceGrpc.ModelServiceBlockingStub modelServiceBlockingStub(ManagedChannel managedChannel) {
        return ModelServiceGrpc.newBlockingStub(managedChannel);
    }

    // 下面的 API 较为底层，暂时不建议使用
    @Bean
    public SessionServiceGrpc.SessionServiceBlockingStub sessionServiceBlockingStub(ManagedChannel managedChannel) {
        return SessionServiceGrpc.newBlockingStub(managedChannel);
    }
//
//    @Bean
//    public MasterServiceGrpc.MasterServiceBlockingStub masterServiceBlockingStub(ManagedChannel managedChannel) {
//        return MasterServiceGrpc.newBlockingStub(managedChannel);
//    }
//
//    @Bean
//    public WorkerServiceGrpc.WorkerServiceBlockingStub workerServiceBlockingStub(ManagedChannel managedChannel) {
//        return WorkerServiceGrpc.newBlockingStub(managedChannel);
//    }
//
//    @Bean
//    public EagerServiceGrpc.EagerServiceBlockingStub eagerServiceBlockingStub(ManagedChannel managedChannel) {
//        return EagerServiceGrpc.newBlockingStub(managedChannel);
//    }
}
